import asyncio
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain_core.prompts import PromptTemplate

from langchain_ollama import ChatOllama
from langchain.output_parsers import MarkdownListOutputParser

from langchain_core.prompts import ChatPromptTemplate
from langchain.agents import AgentExecutor,create_react_agent
from langchain_core.messages import SystemMessage, HumanMessage
#from langchain.agents import tool

#from langchain_core.tools import tool
from langchain.tools import tool

from mcp.server.fastmcp import FastMCP

import json

llm =   ChatOllama(model="qwen2.5:1.5b", temperature=0)
#llm =   ChatOllama(model="deepseek-R1:1.5b", temperature=0)
server_params = StdioServerParameters(
    command="python",
    args=["./mcp_s1.py"],
)

async def run_agent(toolsName,quest):
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            #tools = [multiply,status]
            print("start..........................1")

            tools = await load_mcp_tools(session)
            output_parser = MarkdownListOutputParser()
            print("start..........................2")
            print(tools)
            print("start..........................3")
            #print(tools)
            template = '''Answer the following questions as best you can. You have access to the following tools:
{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!

Question: {input}
Thought:{agent_scratchpad}'''

            prompt = PromptTemplate.from_template(template)
            # 创建并执行agent
            agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)    
            # 创建代理执行器
            agent_executor = AgentExecutor(
                handle_parsing_errors=True,
                agent=agent,
                tools=tools,verbose=True)

           
            await agent_executor.ainvoke({"input":quest})
            return "agent_response"



# Run the async function
if __name__ == "__main__":


    strTools = "请使用 客户设备信息查询工具 工具"
    quest    = " 帮我查询一下 设备名称：试运行设备 的地址是多少，通统计试运行设备的数量 "
    result = asyncio.run(run_agent(strTools,quest))

    print(result)